#include "UdpServer.h"
#include <ws2tcpip.h>
#include <iostream>

UdpServer::UdpServer() : server_fd(INVALID_SOCKET) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        error_msg = "WSAStartup failed";
    }
}

UdpServer::~UdpServer() {
    Close();
}

bool UdpServer::Initialize(const std::string& ip, int port) {
    server_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (server_fd == INVALID_SOCKET) {
        error_msg = "Socket creation failed";
        return false;
    }

    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    inet_pton(AF_INET, ip.c_str(), &server_addr.sin_addr);

    if (bind(server_fd, (sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        error_msg = "Bind failed";
        Close();
        return false;
    }

    return true;
}

bool UdpServer::Send(const std::string& data, const std::string& targetIp, int targetPort) {
    sockaddr_in target_addr{};
    target_addr.sin_family = AF_INET;
    target_addr.sin_port = htons(targetPort);
    inet_pton(AF_INET, targetIp.c_str(), &target_addr.sin_addr);

    int sent = sendto(server_fd, data.c_str(), data.length(), 0, (sockaddr*)&target_addr, sizeof(target_addr));
    if (sent == SOCKET_ERROR) {
        error_msg = "Send failed";
        return false;
    }
    return true;
}

bool UdpServer::Receive(std::function<void(const std::string&, const std::string&, int)> callback) {
    sockaddr_in client_addr{};
    int addr_len = sizeof(client_addr);
    char buffer[1024] = { 0 };

    int received = recvfrom(server_fd, buffer, sizeof(buffer), 0, (sockaddr*)&client_addr, &addr_len);
    if (received == SOCKET_ERROR) {
        error_msg = "Receive failed";
        return false;
    }

    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
    callback(std::string(buffer, received), client_ip, ntohs(client_addr.sin_port));
    return true;
}

void UdpServer::Close() {
    if (server_fd != INVALID_SOCKET) {
        closesocket(server_fd);
        server_fd = INVALID_SOCKET;
    }
    WSACleanup();
}

std::string UdpServer::GetError() const {
    return error_msg;
}